package com.anonymous.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutput;
import java.io.ObjectOutputStream;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Security;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Date;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;

import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Object;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectKeyIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x509.X509CertificateStructure;
import org.bouncycastle.asn1.x509.X509Extension;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.crypto.params.AsymmetricKeyParameter;
import org.bouncycastle.crypto.util.PrivateKeyFactory;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.DefaultDigestAlgorithmIdentifierFinder;
import org.bouncycastle.operator.DefaultSignatureAlgorithmIdentifierFinder;
import org.bouncycastle.operator.bc.BcRSAContentSignerBuilder;

import com.anonymous.message.LogoutMsg;
import com.anonymous.message.SecureMessage;
import com.anonymous.message.SecureObject;

public class ConfidentialitySecurityUtils {

	private static Cipher cipher;	
	/**
	 * A source of secure random numbers
	 */
	static private SecureRandom secureRandom;

	public ConfidentialitySecurityUtils() throws NoSuchAlgorithmException,
			NoSuchPaddingException {
		cipher = Cipher.getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");
	}

	// Generates a key pair using the RSA algorithm
	public KeyPair generateCustomKeyPair() {
		KeyPair pair = null;
		try {
			// Use BC as provider to meet the need of SSL
			Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
			KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA", "BC");
			// Generates random number to be used as a seed
			SecureRandom random = SecureRandom.getInstance("SHA1PRNG", "SUN");
			keyGen.initialize(3072, random);
			pair = keyGen.generateKeyPair();
		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
		return pair;
	}

	// This method encrypts the object with the public key
	// The client calls this method while communicating with the server
	public byte[] encryptMessageObjectAsByteArray(Object obj, PublicKey pub) {
		byte[] c = null;
		try {
			// Serialize to a byte array
			ByteArrayOutputStream bos = new ByteArrayOutputStream();
			ObjectOutputStream out = new ObjectOutputStream(bos);
			out.writeObject(obj);
			out.close();

			// Get the bytes of the serialized object
			byte[] b1 = bos.toByteArray();

			// The encoding function used is OAEP.
			Cipher enc = Cipher
					.getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");

			// Client: c = Enc(m1, K)
			enc.init(Cipher.ENCRYPT_MODE, pub);
			c = enc.doFinal(b1);
		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
		return c;
	}

	public Object encryptMessageObjectAsObject(Object obj, PublicKey pub) {
		Object o = null;
		try {
			// Serialize to a byte array
			ByteArrayOutputStream bos = new ByteArrayOutputStream();
			ObjectOutputStream out = new ObjectOutputStream(bos);
			out.writeObject(obj);
			out.close();
			// Get the bytes of the serialized object
			byte[] b1 = bos.toByteArray();

			System.out.println(b1);
			// The encoding function used is OAEP.
			Cipher enc = Cipher
					.getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");

			// Client: c = Enc(m1, K)
			enc.init(Cipher.ENCRYPT_MODE, pub);
			byte[] c = enc.doFinal(b1);

			// covert to object
			ObjectInputStream in = new ObjectInputStream(
					new ByteArrayInputStream(c));
			o = (Object) in.readObject();
			in.close();
			return o;
		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
		return o;
	}

	// This method decrypts the object with the private key
	// The Server calls this method when
	public Object decryptMessageObjectTakesByteArray(byte[] b1,
			PrivateKey privateKey) {
		Object objOut = null;
		try {

			// Server: m2 = Dec(c; k)
			Cipher dec = Cipher
					.getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");
			dec.init(Cipher.DECRYPT_MODE, privateKey);
			byte[] b2 = dec.doFinal(b1);

			// Deserialize from a byte array
			ObjectInputStream in = new ObjectInputStream(
					new ByteArrayInputStream(b2));
			objOut = (Object) in.readObject();
			in.close();

		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
		return objOut;
	}

	public Object decryptMessageObjectTakesObject(Object objIn,
			PrivateKey privateKey) {
		Object objOut = null;
		try {

			// serialize an object to a byte array
			ByteArrayOutputStream bos = new ByteArrayOutputStream();
			ObjectOutput out = new ObjectOutputStream(bos);
			out.writeObject(objIn);
			out.close();

			byte[] b1 = bos.toByteArray();

			// Server: m2 = Dec(c; k)
			Cipher dec = Cipher
					.getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");
			dec.init(Cipher.DECRYPT_MODE, privateKey);
			byte[] b2 = dec.doFinal(b1);

			// Deserialize from a byte array
			ObjectInputStream in = new ObjectInputStream(
					new ByteArrayInputStream(b2));
			objOut = (Object) in.readObject();
			in.close();

		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
		return objOut;
	}

	public byte[] encrypt(String plaintext, PublicKey pub) throws Exception {
		// public String encrypt(String plaintext, PublicKey pub) throws
		// Exception {
		// Cipher cipher = Cipher
		// .getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");
		cipher.init(Cipher.ENCRYPT_MODE, pub);
		byte[] bytes = plaintext.getBytes("UTF-8");

		byte[] encrypted = blockCipher(bytes, Cipher.ENCRYPT_MODE);

		// char[] encryptedTranspherable = Hex.encodeHex(encrypted);
		// return new String(encryptedTranspherable);
		return encrypted;
	}

	public String decrypt(byte[] b1, PrivateKey privateKey) throws Exception {
		// public String decrypt(String encrypted, PrivateKey privateKey)
		// throws Exception {
		// Cipher cipher = Cipher
		// .getInstance("RSA/ECB/OAEPwithSHA-256andMGF1padding");
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		// byte[] bts = Hex.decodeHex(encrypted.toCharArray());

		byte[] decrypted = blockCipher(b1, Cipher.DECRYPT_MODE);
		// byte[] decrypted = blockCipher(bts, Cipher.DECRYPT_MODE);

		return new String(decrypted, "UTF-8");
	}

	private byte[] blockCipher(byte[] bytes, int mode)
			throws IllegalBlockSizeException, BadPaddingException {
		// string initialize 2 buffers.
		// scrambled will hold intermediate results
		byte[] scrambled = new byte[0];

		// toReturn will hold the total result
		byte[] toReturn = new byte[0];
		// if we encrypt we use 100 byte long blocks. Decryption requires 128
		// byte long blocks (because of RSA)
		int length = (mode == Cipher.ENCRYPT_MODE) ? 100 : 128;

		// another buffer. this one will hold the bytes that have to be modified
		// in this step
		byte[] buffer = new byte[length];

		for (int i = 0; i < bytes.length; i++) {

			// if we filled our buffer array we have our block ready for de- or
			// encryption
			if ((i > 0) && (i % length == 0)) {
				// execute the operation
				scrambled = cipher.doFinal(buffer);
				// add the result to our total result.
				toReturn = append(toReturn, scrambled);
				// here we calculate the length of the next buffer required
				int newlength = length;

				// if newlength would be longer than remaining bytes in the
				// bytes array we shorten it.
				if (i + length > bytes.length) {
					newlength = bytes.length - i;
				}
				// clean the buffer array
				buffer = new byte[newlength];
			}
			// copy byte into our buffer.
			buffer[i % length] = bytes[i];
		}

		// this step is needed if we had a trailing buffer. should only happen
		// when encrypting.
		// example: we encrypt 110 bytes. 100 bytes per run means we "forgot"
		// the last 10 bytes. they are in the buffer array
		scrambled = cipher.doFinal(buffer);

		// final step before we can return the modified data.
		toReturn = append(toReturn, scrambled);

		return toReturn;
	}

	private byte[] append(byte[] prefix, byte[] suffix) {
		byte[] toReturn = new byte[prefix.length + suffix.length];
		for (int i = 0; i < prefix.length; i++) {
			toReturn[i] = prefix[i];
		}
		for (int i = 0; i < suffix.length; i++) {
			toReturn[i + prefix.length] = suffix[i];
		}
		return toReturn;
	}

	public static void main(String[] args) {
		try {
			// test
			// Object objSecurityInEncrypt = "mahitha";
			// System.out.println(objSecurityInEncrypt);

			ConfidentialitySecurityUtils obj = new ConfidentialitySecurityUtils();
			KeyPair pair = obj.generateCustomKeyPair();
			PublicKey publicKey = pair.getPublic();
			PrivateKey privateKey = pair.getPrivate();

			byte[] encryptedObject = obj.encryptMessageObjectAsByteArray(
					new LogoutMsg(), publicKey);
			LogoutMsg logoutMessage = (LogoutMsg) obj
					.decryptMessageObjectTakesByteArray(encryptedObject,
							privateKey);
			System.out.println(logoutMessage.getFlag());
			// byte[] bEncryptOut = null;
			/*
			 * Object objSecurityOutDecrypt = null; ConfidentialitySecurityUtils
			 * securityUtils = new ConfidentialitySecurityUtils(); KeyPair pair
			 * = securityUtils.generateCustomKeyPair();
			 * 
			 * PublicKey pk = pair.getPublic(); PrivateKey priv =
			 * pair.getPrivate();
			 * 
			 * bEncryptOut = securityUtils.encryptMessageObjectAsByteArray(
			 * objSecurityInEncrypt, pair.getPublic());
			 * 
			 * objSecurityOutDecrypt =
			 * securityUtils.decryptMessageObjectTakesByteArray( bEncryptOut,
			 * pair.getPrivate());
			 */

			// test
			// System.out.println(objSecurityOutDecrypt);

			// test
			/*
			 * if (objSecurityOutDecrypt.equals(objSecurityInEncrypt)) {
			 * System.out.println("RSA test Succeeded"); } else {
			 * System.out.println("RSA test Fail"); }
			 */

		} catch (Exception e) {
			e.printStackTrace();
			System.out.println("RSA Test Failed");
		}
	}

	public SecureObject encryptAES(Object obj, SecretKey sessionKey)
			throws Exception {

		Cipher enc = Cipher.getInstance("AES/CBC/PKCS5PADDING");
		enc.init(Cipher.ENCRYPT_MODE, sessionKey);

		ByteArrayOutputStream bos = new ByteArrayOutputStream();
		ObjectOutput out = new ObjectOutputStream(bos);
		out.writeObject(obj);
		out.close();

		byte[] b1 = bos.toByteArray();
		byte[] iv = enc.getIV();
		byte[] secureMessageByteArrayEnc = enc.doFinal(b1);
		return new SecureObject(secureMessageByteArrayEnc, iv);
	}

	public SecureMessage decryptAES(SecureObject secureObject,
			SecretKey sessionKey) throws InvalidKeyException,
			InvalidAlgorithmParameterException, NoSuchAlgorithmException,
			NoSuchPaddingException, IllegalBlockSizeException,
			BadPaddingException, IOException, ClassNotFoundException {
		Cipher dec = Cipher.getInstance("AES/CBC/PKCS5PADDING");
		dec.init(Cipher.DECRYPT_MODE, sessionKey, new IvParameterSpec(
				secureObject.getIv()));
		byte[] SecureMessageByteArray = dec.doFinal(secureObject
				.getSecureMessageByteArrayEnc());

		ObjectInputStream in = new ObjectInputStream(new ByteArrayInputStream(
				SecureMessageByteArray));
		return (SecureMessage) in.readObject();

	}

	public byte[] encryptRSA(byte[] m, PublicKey publicKey)
			throws InvalidKeyException, IllegalBlockSizeException,
			BadPaddingException {
		cipher.init(Cipher.ENCRYPT_MODE, publicKey);
		return cipher.doFinal(m);
	}

	public byte[] decryptRSA(byte[] mEnc, PrivateKey privateKey)
			throws InvalidKeyException, IllegalBlockSizeException,
			BadPaddingException {
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		return cipher.doFinal(mEnc);
	}

	public KeyStore getServerKeystore() throws GeneralSecurityException,
			IOException {
		KeyStore serverKeyStore = KeyStore.getInstance("JKS");

		serverKeyStore.load(new FileInputStream("server.public"),
				"public".toCharArray());
		return serverKeyStore;
	}

	public KeyStore setupClientKeyStore(String passphrase,
			KeyPair clientKeyPair, X509Certificate clientCertificate)
			throws Exception {
		KeyStore clientKeyStore = KeyStore.getInstance("JKS");
		// GENERATE AN EMPTY clientKeyStore
		clientKeyStore.load(null, passphrase.toCharArray());

		X509Certificate clientCert[] = new X509Certificate[1];
		clientCert[0] = clientCertificate;
		KeyStore.PrivateKeyEntry clientPrivateKeyEntry = new KeyStore.PrivateKeyEntry(
				clientKeyPair.getPrivate(), clientCert);
		clientKeyStore.setEntry("clientPrivateKey", clientPrivateKeyEntry,
				new KeyStore.PasswordProtection(passphrase.toCharArray()));

		// Add the certificate
		clientKeyStore.setCertificateEntry("clientCertificate", clientCert[0]);
		return clientKeyStore;
	}
	
	public X509Certificate generateSelfSignedX509Certificate(KeyPair keyPair, String userName)
			throws Exception {

		// yesterday
		Date validityBeginDate = new Date(System.currentTimeMillis() - 24 * 60
				* 60 * 1000);
		// in 2 years
		Date validityEndDate = new Date(System.currentTimeMillis() + 2 * 365
				* 24 * 60 * 60 * 1000);

		// GENERATE THE PUBLIC/PRIVATE RSA KEY PAIR

		// GENERATE THE X509 CERTIFICATE
		SubjectPublicKeyInfo subjectPublicKeyInfo = new SubjectPublicKeyInfo(
				(ASN1Sequence) ASN1Object.fromByteArray(keyPair.getPublic()
						.getEncoded()));
		X500Name dnName = new X500Name("CN=" + userName);

		X509v3CertificateBuilder certBuilder = new X509v3CertificateBuilder(
				dnName, BigInteger.valueOf(System.currentTimeMillis()),
				validityBeginDate, validityEndDate, dnName,
				subjectPublicKeyInfo);

		certBuilder.addExtension(X509Extension.subjectKeyIdentifier, false,
				createSubjectKeyId(keyPair.getPublic()));

		// prepare the signer with the private Key
		AlgorithmIdentifier sigAlgId = new DefaultSignatureAlgorithmIdentifierFinder()
				.find("SHA256withRSA");
		AlgorithmIdentifier digAlgId = new DefaultDigestAlgorithmIdentifierFinder()
				.find(sigAlgId);

		BcRSAContentSignerBuilder sigBuild = new BcRSAContentSignerBuilder(
				sigAlgId, digAlgId);
		
		// hopefully format is PKCS#8
		AsymmetricKeyParameter foo = PrivateKeyFactory.createKey(keyPair
				.getPrivate().getEncoded());

		ContentSigner sigGen = sigBuild.build(foo);

		X509CertificateHolder certHolder = certBuilder.build(sigGen);
		X509CertificateStructure eeX509CertificateStructure = certHolder
				.toASN1Structure();

		// retrieve the certificate from holder

		CertificateFactory cf = CertificateFactory.getInstance("X.509");

		// Read Certificate
		InputStream is1 = new ByteArrayInputStream(
				eeX509CertificateStructure.getEncoded());
		X509Certificate theCert = (X509Certificate) cf.generateCertificate(is1);
		is1.close();

		return theCert;
	}

	public static SubjectKeyIdentifier createSubjectKeyId(PublicKey pubKey) {
		try {
			ByteArrayInputStream bIn = new ByteArrayInputStream(
					pubKey.getEncoded());
			SubjectPublicKeyInfo info = new SubjectPublicKeyInfo(
					(ASN1Sequence) new ASN1InputStream(bIn).readObject());

			return new SubjectKeyIdentifier(info);
		} catch (Exception e) {
			throw new RuntimeException("ERROR: Creating Subject Key ID");
		}
	}
	
	public SSLContext setupSSLContext(String passphrase, KeyStore clientKeyStore, KeyStore serverKeyStore) throws Exception{
		TrustManagerFactory tmf = TrustManagerFactory.getInstance("SunX509");
		tmf.init(clientKeyStore);

		KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
		kmf.init(serverKeyStore, passphrase.toCharArray());

		SSLContext sslContext = SSLContext.getInstance("TLS");
		

		secureRandom = new SecureRandom();
		secureRandom.nextInt();
		sslContext.init(kmf.getKeyManagers(), tmf.getTrustManagers(),
				secureRandom);
		
		return sslContext;
	}

}
